import os
import base64
from pathlib import Path
from PIL import Image
import numpy as np
import cv2
    
from Chinese_Chess_Recognition.core.chessboard_detector import ChessboardDetector

detector = ChessboardDetector(
    pose_model_path="onnx/pose/4_v6-0301.onnx",
    full_classifier_model_path="onnx/layout_recognition/nano_v3-0319.onnx"
)

# 数据集路径
dict_cate_names = {
    '.': '.',
    'x': 'x',
    '红帅': 'K',
    '红士': 'A',
    '红相': 'B',
    '红马': 'N',
    '红车': 'R',
    '红炮': 'C',
    '红兵': 'P',

    '黑将': 'k',
    '黑仕': 'a',
    '黑象': 'b',
    '黑傌': 'n',
    '黑車': 'r',
    '黑砲': 'c',
    '黑卒': 'p',
}

# 数据集路径
dict_cate_images = {
    'K': 'red_K.png',
    'A': 'red_A.png',
    'B': 'red_B.png',
    'N': 'red_N.png',
    'R': 'red_R.png',
    'C': 'red_C.png',
    'P': 'red_P.png',
    'k': 'black_k.png',
    'a': 'black_a.png',
    'b': 'black_b.png',
    'n': 'black_n.png',
    'r': 'black_r.png',
    'c': 'black_c.png',
    'p': 'black_p.png',
}

dict_cate_names_reverse = {v: k for k, v in dict_cate_names.items()}

def detect_chessboard(image):
    if image is None:
        return None, None, None, None, None
    
    try:
        original_image_with_keypoints, transformed_image, cells_labels_str, scores, time_info = detector.pred_detect_board_and_classifier(image)

        # 将 cells_labels 转换为 DataFrame
        # cells_labels 通过  \n 分割
        annotation_10_rows = [item for item in cells_labels_str.split("\n")]
        # 将 annotation_10_rows 转换成为 10 行 9 列的二维数组
        annotation_arr_10_9_short = [list(item) for item in annotation_10_rows]

        # 将 棋子类别 转换为 中文
        annotation_arr_10_9 = [[dict_cate_names_reverse[item] for item in row] for row in annotation_arr_10_9_short]
        
        return annotation_arr_10_9

    except Exception as e:
        print(f"检测失败 图片或者视频布局错误")
        return None, None, None, None, None
    
def get_all_file_paths(folder_path, extensions=None):
    """
    获取文件夹中所有图像文件的绝对路径

    Args:
        folder_path (str): 文件夹路径
        extensions (list): 文件扩展名列表，如 ['.png', '.jpg']，None表示所有文件

    Returns:
        list: 所有文件的绝对路径列表
    """
    file_paths = []
    folder = Path(folder_path)

    if not folder.exists():
        print(f"文件夹不存在: {folder_path}")
        return file_paths

    if extensions is None:
        extensions = ['.png', '.jpg', '.jpeg', '.bmp', '.tiff', '.tif']

    for file_path in folder.rglob('*'):
        if file_path.is_file():
            if extensions is None or file_path.suffix.lower() in extensions:
                file_paths.append(str(file_path.absolute()))

    return sorted(file_paths)

paths = get_all_file_paths(folder_path="/home/knowin/codes/knowin-robot/boards")

for image in paths:
    print(image)
    image = Image.open(image)
    arr = np.array(image)
    res = detect_chessboard(arr)
    print(res)
